"""Metrics collector for monitoring and analysis."""

from typing import Dict, Any, Optional, List
from dataclasses import dataclass, field
from datetime import datetime, timedelta
import logging
import json
from collections import defaultdict

logger = logging.getLogger(__name__)


@dataclass
class ModelMetrics:
    """Metrics for a specific model."""
    
    model_name: str
    total_requests: int = 0
    successful_requests: int = 0
    failed_requests: int = 0
    total_tokens: int = 0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_cost: float = 0.0
    total_latency: float = 0.0
    min_latency: Optional[float] = None
    max_latency: Optional[float] = None
    avg_latency: float = 0.0
    last_used: Optional[datetime] = None
    error_types: Dict[str, int] = field(default_factory=dict)
    
    def record_request(
        self,
        success: bool,
        latency: float,
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
        cost: float = 0.0,
        error_type: Optional[str] = None
    ):
        """Record a request.
        
        Args:
            success: Whether request succeeded
            latency: Request latency
            prompt_tokens: Prompt tokens used
            completion_tokens: Completion tokens generated
            cost: Request cost
            error_type: Type of error if failed
        """
        self.total_requests += 1
        
        if success:
            self.successful_requests += 1
        else:
            self.failed_requests += 1
            if error_type:
                self.error_types[error_type] = self.error_types.get(error_type, 0) + 1
        
        self.total_latency += latency
        self.avg_latency = self.total_latency / self.total_requests
        
        if self.min_latency is None or latency < self.min_latency:
            self.min_latency = latency
        if self.max_latency is None or latency > self.max_latency:
            self.max_latency = latency
        
        self.prompt_tokens += prompt_tokens
        self.completion_tokens += completion_tokens
        self.total_tokens += prompt_tokens + completion_tokens
        self.total_cost += cost
        
        self.last_used = datetime.utcnow()
    
    def get_success_rate(self) -> float:
        """Get success rate.
        
        Returns:
            Success rate as percentage
        """
        if self.total_requests == 0:
            return 0.0
        return (self.successful_requests / self.total_requests) * 100
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary.
        
        Returns:
            Metrics as dictionary
        """
        return {
            "model_name": self.model_name,
            "total_requests": self.total_requests,
            "successful_requests": self.successful_requests,
            "failed_requests": self.failed_requests,
            "success_rate": self.get_success_rate(),
            "total_tokens": self.total_tokens,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_cost": self.total_cost,
            "avg_latency": self.avg_latency,
            "min_latency": self.min_latency,
            "max_latency": self.max_latency,
            "last_used": self.last_used.isoformat() if self.last_used else None,
            "error_types": self.error_types
        }


class MetricsCollector:
    """Collector for aggregating metrics."""
    
    def __init__(self, window_size_minutes: int = 60):
        """Initialize metrics collector.
        
        Args:
            window_size_minutes: Time window for rolling metrics
        """
        self.window_size = timedelta(minutes=window_size_minutes)
        self.model_metrics: Dict[str, ModelMetrics] = {}
        self.chain_metrics: Dict[str, Dict[str, Any]] = defaultdict(lambda: {
            "executions": 0,
            "successes": 0,
            "failures": 0,
            "total_duration": 0.0,
            "avg_duration": 0.0
        })
        self.time_series: List[Dict[str, Any]] = []
        self.max_time_series_size = 1000
    
    def record_llm_request(
        self,
        model: str,
        success: bool,
        latency: float,
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
        cost: float = 0.0,
        error_type: Optional[str] = None
    ):
        """Record an LLM request.
        
        Args:
            model: Model name
            success: Whether request succeeded
            latency: Request latency
            prompt_tokens: Prompt tokens
            completion_tokens: Completion tokens
            cost: Request cost
            error_type: Error type if failed
        """
        # Update model metrics
        if model not in self.model_metrics:
            self.model_metrics[model] = ModelMetrics(model_name=model)
        
        self.model_metrics[model].record_request(
            success=success,
            latency=latency,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            cost=cost,
            error_type=error_type
        )
        
        # Record time series data
        self._record_time_series({
            "timestamp": datetime.utcnow().isoformat(),
            "type": "llm_request",
            "model": model,
            "success": success,
            "latency": latency,
            "tokens": prompt_tokens + completion_tokens,
            "cost": cost
        })
    
    def record_chain_execution(
        self,
        chain_type: str,
        success: bool,
        duration: float
    ):
        """Record a chain execution.
        
        Args:
            chain_type: Type of chain
            success: Whether execution succeeded
            duration: Execution duration
        """
        metrics = self.chain_metrics[chain_type]
        metrics["executions"] += 1
        
        if success:
            metrics["successes"] += 1
        else:
            metrics["failures"] += 1
        
        metrics["total_duration"] += duration
        metrics["avg_duration"] = metrics["total_duration"] / metrics["executions"]
        
        # Record time series
        self._record_time_series({
            "timestamp": datetime.utcnow().isoformat(),
            "type": "chain_execution",
            "chain_type": chain_type,
            "success": success,
            "duration": duration
        })
    
    def _record_time_series(self, data: Dict[str, Any]):
        """Record time series data point.
        
        Args:
            data: Data point to record
        """
        self.time_series.append(data)
        
        # Limit size
        if len(self.time_series) > self.max_time_series_size:
            self.time_series = self.time_series[-self.max_time_series_size:]
    
    def get_model_metrics(self, model: Optional[str] = None) -> Dict[str, Any]:
        """Get model metrics.
        
        Args:
            model: Specific model or None for all
            
        Returns:
            Metrics dictionary
        """
        if model:
            if model in self.model_metrics:
                return self.model_metrics[model].to_dict()
            return {}
        
        return {
            name: metrics.to_dict()
            for name, metrics in self.model_metrics.items()
        }
    
    def get_chain_metrics(self, chain_type: Optional[str] = None) -> Dict[str, Any]:
        """Get chain metrics.
        
        Args:
            chain_type: Specific chain or None for all
            
        Returns:
            Metrics dictionary
        """
        if chain_type:
            return dict(self.chain_metrics.get(chain_type, {}))
        
        return dict(self.chain_metrics)
    
    def get_summary(self) -> Dict[str, Any]:
        """Get metrics summary.
        
        Returns:
            Summary dictionary
        """
        total_requests = sum(m.total_requests for m in self.model_metrics.values())
        total_tokens = sum(m.total_tokens for m in self.model_metrics.values())
        total_cost = sum(m.total_cost for m in self.model_metrics.values())
        
        return {
            "total_requests": total_requests,
            "total_tokens": total_tokens,
            "total_cost": total_cost,
            "models": len(self.model_metrics),
            "chains": len(self.chain_metrics),
            "time_series_points": len(self.time_series),
            "model_metrics": self.get_model_metrics(),
            "chain_metrics": self.get_chain_metrics()
        }
    
    def get_recent_metrics(
        self,
        minutes: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """Get recent time series metrics.
        
        Args:
            minutes: Number of minutes to look back
            
        Returns:
            List of recent metrics
        """
        if not minutes:
            minutes = self.window_size.total_seconds() / 60
        
        cutoff = datetime.utcnow() - timedelta(minutes=minutes)
        
        recent = []
        for point in self.time_series:
            timestamp = datetime.fromisoformat(point["timestamp"])
            if timestamp >= cutoff:
                recent.append(point)
        
        return recent
    
    def export_metrics(self) -> str:
        """Export metrics as JSON.
        
        Returns:
            JSON string of metrics
        """
        return json.dumps(self.get_summary(), indent=2, default=str)
    
    def reset(self):
        """Reset all metrics."""
        self.model_metrics.clear()
        self.chain_metrics.clear()
        self.time_series.clear()
        logger.info("Metrics collector reset")